module Zeno.Type where

import Prelude ()
import StdImports
import qualified Data.Set as Set

data Type a 
  = TypeVar !a
  | FunType !(Type a) !(Type a)
  | AppType !(Type a) !(Type a)
  | EmptyType
  deriving ( Eq, Ord, Functor, Foldable, Traversable )

instance NFData (Type a) 

instance WithinTraversable (Type a) (Type a) where
  mapWithinM f (FunType t1 t2) = 
    f =<< FunType `liftM` mapWithinM f t1 `ap` mapWithinM f t2
  mapWithinM f (AppType t1 t2) = 
    f =<< AppType `liftM` mapWithinM f t1 `ap` mapWithinM f t2
  mapWithinM f t = f t
  
type TypeSubstitution a = Substitution (Type a)

instance Ord a => Unifiable (Type a) where
  unify t1 t2 = 
    case (t1, t2) of
      (FunType {}, FunType {}) -> fromMaybe backup $ 
        unifyFlattened (flattenFunType t1) (flattenFunType t2)
      (AppType {}, AppType {}) -> fromMaybe backup $
        unifyFlattened (flattenAppType t1) (flattenAppType t2)
      _ -> backup
    where
      backup = Set.singleton (t1, t2)
      
      unifyFlattened (f1 : as1) (f2 : as2) = 
        if f1 == f2 && length as1 == length as2 
          then Just $ mconcat (zipWith unify as1 as2)
          else Nothing

unflattenFunType :: [Type a] -> Type a
unflattenFunType [] = EmptyType
unflattenFunType ts = foldr1 FunType ts

unflattenAppType :: [Type a] -> Type a
unflattenAppType [] = EmptyType
unflattenAppType ts = foldl1 AppType ts

fromTypeVar :: Type a -> a
fromTypeVar (TypeVar a) = a

isTypeVar :: Type a -> Bool
isTypeVar (TypeVar _) = True
isTypeVar _ = False 

flattenFunType :: Type a -> [Type a]
flattenFunType (FunType f a) = f : flattenFunType a
flattenFunType t = [t]

flattenAppType :: Type a -> [Type a]
flattenAppType (AppType f a) = flattenAppType f ++ [a]
flattenAppType t = [t]

isFunctionType :: Type a -> Bool
isFunctionType (FunType _ _) = True
isFunctionType _ = False

isEmptyType :: Type a -> Bool
isEmptyType EmptyType = True
isEmptyType _ = False

instance Show a => Show (Type a) where
  show EmptyType = "_|_"
  show (AppType fun arg) = 
    show fun ++ " " ++ show_arg
    where
      show_arg = 
        case arg of
          AppType  _ _ -> "(" ++ show arg ++ ")"
          _ -> show arg
  show (TypeVar id) = show id
  show (FunType arg res) = 
    show_arg ++ " -> " ++ show res
    where 
      show_arg = 
        case arg of 
          FunType _ _ -> "(" ++ show arg ++ ")"
          _ -> show arg
